#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
梯形加速度轨迹规划器
"""

import numpy as np
import casadi as ca
from diff_robot_model import build_robot_model

class TrapezoidalTrajectoryPlanner:
    def __init__(self, max_acceleration=1.0, max_velocity=0.6):
        """
        初始化梯形轨迹规划器
        
        Args:
            max_acceleration: 最大加速度 (m/s^2)
            max_velocity: 最大速度 (m/s)
        """
        self.max_acceleration = max_acceleration
        self.max_velocity = max_velocity
        
    def plan_linear_trajectory(self, start_pos, end_pos, duration):
        """
        规划直线轨迹
        
        Args:
            start_pos: 起始位置 [x, y]
            end_pos: 终点位置 [x, y]
            duration: 总时间 (s)
            
        Returns:
            time_points: 时间点数组
            positions: 位置数组
            velocities: 速度数组
            accelerations: 加速度数组
        """
        # 计算总距离和方向
        distance = np.linalg.norm(np.array(end_pos) - np.array(start_pos))
        direction = (np.array(end_pos) - np.array(start_pos)) / distance if distance > 1e-6 else np.array([1.0, 0.0])
        
        # 计算梯形速度曲线的时间段
        # 加速时间
        t_acc = self.max_velocity / self.max_acceleration
        # 加速阶段行驶距离
        d_acc = 0.5 * self.max_velocity * t_acc
        # 匀速阶段行驶距离
        d_const = distance - 2 * d_acc
        
        if d_const >= 0:
            # 典型梯形速度曲线
            t_const = d_const / self.max_velocity
            total_time = 2 * t_acc + t_const
            # 调整时间以匹配指定的duration
            time_scale = duration / total_time
            t_acc_scaled = t_acc * time_scale
            t_const_scaled = t_const * time_scale
        else:
            # 三角形速度曲线（最大速度达不到）
            max_vel_possible = np.sqrt(distance * self.max_acceleration)
            t_acc_scaled = np.sqrt(distance / self.max_acceleration)
            t_const_scaled = 0
            time_scale = duration / (2 * t_acc_scaled)
            self.max_velocity = max_vel_possible
        
        # 生成时间点
        dt = 0.1  # 采样时间
        time_points = np.arange(0, duration, dt)
        
        # 计算每个时间点的速度、位置和加速度
        positions = []
        velocities = []
        accelerations = []
        
        for t in time_points:
            if t <= t_acc_scaled:
                # 加速阶段
                acc = self.max_acceleration * time_scale
                vel = acc * t
                pos = 0.5 * acc * t**2
                accel = acc
            elif t <= t_acc_scaled + t_const_scaled:
                # 匀速阶段
                vel = self.max_velocity
                pos = 0.5 * self.max_acceleration * time_scale * t_acc_scaled**2 + vel * (t - t_acc_scaled)
                accel = 0
            elif t <= duration:
                # 减速阶段
                t_dec = t - (t_acc_scaled + t_const_scaled)
                acc = -self.max_acceleration * time_scale
                vel = self.max_velocity + acc * t_dec
                pos = (0.5 * self.max_acceleration * time_scale * t_acc_scaled**2 + 
                       self.max_velocity * t_const_scaled + 
                       self.max_velocity * t_dec + 0.5 * acc * t_dec**2)
                accel = acc
            else:
                # 结束
                vel = 0
                pos = distance
                accel = 0
            
            # 转换为全局坐标
            global_pos = np.array(start_pos) + pos * direction
            global_vel = vel * direction
            global_acc = accel * direction
            
            positions.append(global_pos)
            velocities.append(global_vel)
            accelerations.append(global_acc)
        
        return time_points, np.array(positions), np.array(velocities), np.array(accelerations)

def simulate_robot_trajectory(robot_model, positions, velocities, time_points):
    """
    使用差分小车模型模拟轨迹
    
    Args:
        robot_model: 差分小车模型
        positions: 位置轨迹
        velocities: 速度轨迹
        time_points: 时间点
        
    Returns:
        robot_states: 机器人状态历史
        robot_controls: 机器人控制输入历史
    """
    # 初始化
    robot_states = []
    robot_controls = []
    
    # 初始状态 [x, y, vx, vy, theta]
    current_state = np.array([positions[0][0], positions[0][1], 0.0, 0.0, 0.0])
    
    for i in range(len(time_points)-1):
        # 计算期望的加速度
        if i < len(velocities) - 1:
            desired_acceleration = (velocities[i+1] - velocities[i]) / (time_points[i+1] - time_points[i])
        else:
            desired_acceleration = np.array([0.0, 0.0])
        
        # 计算控制输入
        # 对于差分小车，我们需要将加速度转换为控制输入
        # a = sqrt(ax^2 + ay^2)，omega = (vy*ax - vx*ay) / (vx^2 + vy^2)
        vx = current_state[2]
        vy = current_state[3]
        
        # 计算线加速度大小
        a = np.linalg.norm(desired_acceleration)
        
        # 计算角速度
        if vx**2 + vy**2 > 1e-6:
            omega = (vy * desired_acceleration[0] - vx * desired_acceleration[1]) / (vx**2 + vy**2)
        else:
            omega = 0.0
            
        # 限制控制输入
        a = np.clip(a, -robot_model.a_max, robot_model.a_max)
        omega = np.clip(omega, -robot_model.omega_max, robot_model.omega_max)
        
        control = np.array([a, omega])
        
        # 使用模型计算下一个状态
        state_dm = ca.DM(current_state)
        control_dm = ca.DM(control)
        rhs = robot_model.f(state_dm, control_dm)
        next_state = current_state + robot_model.T * np.array(rhs.full()).flatten()
        
        # 保存状态和控制
        robot_states.append(current_state.copy())
        robot_controls.append(control.copy())
        
        # 更新当前状态
        current_state = next_state
        
        # 确保方向与速度方向一致
        if current_state[2]**2 + current_state[3]**2 > 1e-6:
            current_state[4] = np.arctan2(current_state[3], current_state[2])
    
    robot_states.append(current_state)
    return robot_states, robot_controls